import time
from typing import List
import numpy as np
from tqdm import tqdm
from openai import OpenAI
import openai
import sys
import os

# 添加项目根目录到Python路径
current_dir = os.path.dirname(os.path.abspath(__file__))
project_root = os.path.abspath(os.path.join(current_dir, '..'))
if project_root not in sys.path:
    sys.path.insert(0, project_root)
# =============================================

try:
    from config import settings
    from utils.logger import setup_logger
except ImportError as e:
    print(f"导入错误: {e}")
    print("正在使用备用方案...")


    # 备用方案：直接在这里定义必要的类和函数
    class Settings:
        OPENAI_MODEL = "BAAI/bge-large-zh-v1.5"
        OPENAI_API_KEY = "sk-aswwzdvvimeybiiokqebixpkhbmcftlbgkubssfuodifqjcf"
        OPENAI_BASE_URL = "https://api.siliconflow.cn/v1"
        MAX_BATCH_SIZE = 32
        MAX_RETRY_ATTEMPTS = 5


    settings = Settings()


    def setup_logger(name):
        import logging
        logger = logging.getLogger(name)
        if not logger.handlers:
            handler = logging.StreamHandler()
            formatter = logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s')
            handler.setFormatter(formatter)
            logger.addHandler(handler)
            logger.setLevel(logging.INFO)
        return logger

logger = setup_logger(__name__)


class EmbeddingClient:
    """嵌入向量生成客户端"""

    def __init__(self):
        try:
            self.client = OpenAI(
                base_url=settings.OPENAI_BASE_URL,
                api_key=settings.OPENAI_API_KEY
            )
            self.model = settings.OPENAI_MODEL
            self.max_batch = settings.MAX_BATCH_SIZE
            self.max_retry = settings.MAX_RETRY_ATTEMPTS
            logger.info("EmbeddingClient 初始化成功")
        except Exception as e:
            logger.error(f"EmbeddingClient 初始化失败: {e}")
            # 备用初始化
            self.client = None
            self.model = settings.OPENAI_MODEL
            self.max_batch = settings.MAX_BATCH_SIZE
            self.max_retry = settings.MAX_RETRY_ATTEMPTS

    def create_embeddings(self, texts: List[str], model: str = None) -> np.ndarray:
        """
        为文本列表创建嵌入向量

        Args:
            texts: 文本列表
            model: 模型名称，默认为配置中的模型

        Returns:
            np.ndarray: 嵌入向量数组
        """
        if model is None:
            model = self.model

        # 清洗文本
        cleaned = [(t or "未知") if isinstance(t, str) else "未知" for t in texts]

        # 如果没有配置API客户端，返回随机向量（用于测试）
        if self.client is None:
            logger.warning("使用模拟嵌入向量（测试模式）")
            return np.random.randn(len(cleaned), 768)

        total_batches = (len(cleaned) + self.max_batch - 1) // self.max_batch
        vecs = []

        for batch_idx in tqdm(range(total_batches), desc="生成嵌入向量"):
            batch = cleaned[batch_idx * self.max_batch: (batch_idx + 1) * self.max_batch]

            for attempt in range(self.max_retry):
                try:
                    resp = self.client.embeddings.create(model=model, input=batch)
                    vecs.extend([d.embedding for d in resp.data])
                    break
                except openai.BadRequestError as e:
                    if "20015" in str(e):  # 限流错误
                        wait = 2 ** attempt
                        logger.warning(f"限流重试 {attempt + 1}/{self.max_retry}，等待 {wait}s")
                        time.sleep(wait)
                    else:
                        raise
            else:
                raise RuntimeError(f"批次 {batch_idx} 连续失败 {self.max_retry} 次")

        return np.array(vecs, dtype=np.float32)

    def test_connection(self) -> bool:
        """测试API连接"""
        try:
            test_texts = ["测试连接"]
            embeddings = self.create_embeddings(test_texts)
            success = embeddings.shape[0] == 1
            if success:
                logger.info("API连接测试成功")
            else:
                logger.warning("API连接测试返回异常")
            return success
        except Exception as e:
            logger.error(f"API连接测试失败: {e}")
            return False


# 测试代码
if __name__ == "__main__":
    print("测试 EmbeddingClient...")

    client = EmbeddingClient()

    if client.test_connection():
        print("✓ API连接测试成功")

        # 测试生成嵌入向量
        test_texts = ["你好", "世界", "测试文本"]
        embeddings = client.create_embeddings(test_texts)
        print(f"✓ 成功生成 {len(test_texts)} 个文本的嵌入向量")
        print(f"✓ 向量维度: {embeddings.shape}")
    else:
        print("✗ API连接测试失败，使用模拟模式")

        # 即使失败也测试模拟模式
        test_texts = ["你好", "世界", "测试文本"]
        embeddings = client.create_embeddings(test_texts)
        print(f"✓ 模拟模式生成 {len(test_texts)} 个文本的嵌入向量")
        print(f"✓ 向量维度: {embeddings.shape}")